package backend.opt;

import backend.Instruction.*;
import backend.component.RISCVblock;
import backend.component.RISCVfunction;
import backend.component.RISCVmodule;

import java.util.*;

/**
 * BranchFolding 的目标是合并重复的跳转目标，消除冗余的跳转指令，优化控制流。
 */
public class branchFolding {

    private RISCVmodule module;

    public branchFolding(RISCVmodule riscvmodule) {
        this.module = riscvmodule;
    }

    public void optmize() {
        boolean changed = true;
        // 重复执行直到没有更多优化机会
        while (changed) {
            changed = false;
            //# 多个块都跳转到 ifentry7
            //iffalse6:
            //	li s3, 100
            //	j ifentry7
            //
            //iffalse10:
            //	li s3, 99
            //	j ifentry7
            //
            //iftrue9:
            //	li s3, 1
            //	j ifentry7
            changed |= eliminateRedundantJumps();     // 消除冗余跳转
            //# 复杂的条件分支链
            //and_true8:
            //	slti t0, s3, 100
            //	beqz t0, iffalse6
            //	j iftrue5
            //
            //iftrue5:
            //	beqz s4, iffalse10
            //	j and_true12
            changed |= foldConditionalBranches();     // 折叠条件分支
            //changed |= mergeIdenticalBlocks();        // 合并相同的基本块
        }
    }

    /**
     * 消除冗余跳转：
     * 场景1: j A; A: j B → j B
     * 场景2: beqz r, A; j A → beqz r, A (删除多余的j)
     */
    private boolean eliminateRedundantJumps() {
        boolean changed = false;

        for (RISCVfunction function : module.getFunctions()) {
            for (RISCVblock block : function.getBlocks()) {
                LinkedList<RISCVinst> instructions = block.getInstructionlist();

                // 检查块的最后一条指令
                if (instructions.isEmpty()) continue;
                RISCVinst lastInst = instructions.getLast();

                // 场景1：处理无条件跳转链
                if (lastInst instanceof RISCVj jumpInst) {
                    RISCVblock target = jumpInst.getTarget();
                    RISCVblock finalTarget = resolveJumpChain(target);

                    if (!target.equals(finalTarget)) {
                        jumpInst.setTarget(finalTarget);
                        changed = true;
                    }
                }

                // 场景2：条件跳转后的冗余无条件跳转
                if (instructions.size() >= 2) {
                    RISCVinst secondLast = instructions.get(instructions.size() - 2);

                    if (secondLast instanceof RISCVbeqz branchInst &&
                            lastInst instanceof RISCVj jumpInst) {

                        // 如果条件跳转和无条件跳转指向同一目标，删除无条件跳转
                        if (branchInst.getTarget().equals(jumpInst.getTarget())) {
                            block.deleteInstr(jumpInst);
                            changed = true;
                        }
                        // 或者检查是否可以简化为单个条件跳转
                        else if (canSimplifyBranchJump(branchInst, jumpInst)) {
                            simplifyBranchJump(block, branchInst, jumpInst);
                            changed = true;
                        }
                    }
                }
            }
        }

        return changed;
    }

    /**
     * 解析跳转链，找到最终的跳转目标
     */
    private RISCVblock resolveJumpChain(RISCVblock start) {
        Set<RISCVblock> visited = new HashSet<>();
        RISCVblock current = start;

        while (current != null && !visited.contains(current)) {
            visited.add(current);

            LinkedList<RISCVinst> instructions = current.getInstructionlist();
            if (instructions.size() == 1 && instructions.getFirst() instanceof RISCVj jumpInst) {
                current = jumpInst.getTarget();
            } else {
                break;
            }
        }

        return current != null ? current : start;
    }

    /**
     * 检查是否可以简化 beqz + j 的组合
     */
    private boolean canSimplifyBranchJump(RISCVbeqz branchInst, RISCVj jumpInst) {
        // 检查跳转目标是否可以合并或优化
        RISCVblock branchTarget = branchInst.getTarget();
        RISCVblock jumpTarget = jumpInst.getTarget();

        // 如果分支目标只是一个简单的跳转，可以直接跳转到最终目标
        if (branchTarget.getInstructionlist().size() == 1 &&
                branchTarget.getInstructionlist().getFirst() instanceof RISCVj) {
            return true;
        }

        return false;
    }

    /**
     * 简化 beqz + j 组合
     */
    private void simplifyBranchJump(RISCVblock block, RISCVbeqz branchInst, RISCVj jumpInst) {
        RISCVblock branchTarget = branchInst.getTarget();

        // 如果分支目标是简单跳转，直接跳转到最终目标
        if (branchTarget.getInstructionlist().size() == 1 &&
                branchTarget.getInstructionlist().getFirst() instanceof RISCVj targetJump) {
            branchInst.setTarget(targetJump.getTarget());
        }
    }

    /**
     * 折叠条件分支：
     * 将复杂的条件跳转模式简化
     */
    private boolean foldConditionalBranches() {
        boolean changed = false;

        for (RISCVfunction function : module.getFunctions()) {
            for (RISCVblock block : function.getBlocks()) {
                LinkedList<RISCVinst> instructions = block.getInstructionlist();

                // 寻找可以折叠的条件分支模式
                for (int i = 0; i < instructions.size() - 1; i++) {
                    RISCVinst current = instructions.get(i);
                    RISCVinst next = instructions.get(i + 1);

                    // 模式：连续的条件跳转到相同目标
                    if (current instanceof RISCVbeqz branch1 &&
                            next instanceof RISCVbeqz branch2 &&
                            branch1.getTarget().equals(branch2.getTarget())) {

                        // 可以考虑合并条件（如果逻辑允许）
                        if (canMergeConditions(branch1, branch2)) {
                            mergeConditions(block, branch1, branch2, i);
                            changed = true;
                            break; // 重新开始扫描
                        }
                    }
                }
            }
        }

        return changed;
    }

    /**
     * 检查是否可以合并两个条件
     */
    private boolean canMergeConditions(RISCVbeqz branch1, RISCVbeqz branch2) {
        // 简单情况：如果两个条件使用相同的寄存器
        return branch1.getSrc().equals(branch2.getSrc());
    }

    /**
     * 合并条件分支
     */
    private void mergeConditions(RISCVblock block, RISCVbeqz branch1, RISCVbeqz branch2, int index) {
        // 这里可以根据具体的逻辑关系进行优化
        // 简单实现：保留第一个分支，删除第二个
        block.deleteInstr(branch2);
    }

    /**
     * 合并相同的基本块：
     * 如果多个基本块有相同的指令序列，可以合并它们
     */
    private boolean mergeIdenticalBlocks() {
        boolean changed = false;
        Map<String, List<RISCVblock>> blocksBySignature = new HashMap<>();

        // 为每个基本块计算签名
        for (RISCVfunction function : module.getFunctions()) {
            for (RISCVblock block : function.getBlocks()) {
                String signature = computeBlockSignature(block);
                blocksBySignature.computeIfAbsent(signature, k -> new ArrayList<>()).add(block);
            }
        }

        // 合并具有相同签名的块
        for (List<RISCVblock> identicalBlocks : blocksBySignature.values()) {
            if (identicalBlocks.size() > 1) {
                changed |= mergeBlocks(identicalBlocks);
            }
        }

        return changed;
    }

    /**
     * 计算基本块的签名（用于识别相同的块）
     */
    private String computeBlockSignature(RISCVblock block) {
        StringBuilder signature = new StringBuilder();

        for (RISCVinst inst : block.getInstructionlist()) {
            // 简化的签名计算，只考虑指令类型和操作
            if (inst instanceof RISCVbinary binary) {
                signature.append("BIN:").append(binary.getType()).append(";");
            } else if (inst instanceof RISCVmove) {
                signature.append("MOV;");
            } else if (inst instanceof RISCVj) {
                signature.append("JMP;");
            } else if (inst instanceof RISCVbeqz) {
                signature.append("BEQ;");
            } else {
                signature.append(inst.getClass().getSimpleName()).append(";");
            }
        }

        return signature.toString();
    }

    /**
     * 合并相同的基本块
     */
    private boolean mergeBlocks(List<RISCVblock> blocks) {
        if (blocks.size() < 2) return false;

        RISCVblock keepBlock = blocks.get(0);
        boolean changed = false;

        // 将其他块的引用都指向保留的块
        for (int i = 1; i < blocks.size(); i++) {
            RISCVblock redundantBlock = blocks.get(i);

            // 更新所有指向冗余块的跳转
            updateReferencesToBlock(redundantBlock, keepBlock);

            // 从函数中删除冗余块
            removeBlockFromFunction(redundantBlock);
            changed = true;
        }

        return changed;
    }

    /**
     * 更新所有指向指定块的引用
     */
    private void updateReferencesToBlock(RISCVblock oldBlock, RISCVblock newBlock) {
        for (RISCVfunction function : module.getFunctions()) {
            for (RISCVblock block : function.getBlocks()) {
                for (RISCVinst inst : block.getInstructionlist()) {
                    if (inst instanceof RISCVj jumpInst && jumpInst.getTarget().equals(oldBlock)) {
                        jumpInst.setTarget(newBlock);
                    }
                    if (inst instanceof RISCVbeqz branchInst && branchInst.getTarget().equals(oldBlock)) {
                        branchInst.setTarget(newBlock);
                    }
                }
            }
        }
    }

    /**
     * 从函数中删除指定的基本块
     */
    private void removeBlockFromFunction(RISCVblock block) {
        for (RISCVfunction function : module.getFunctions()) {
            if (function.getBlocks().contains(block)) {
                function.deleteBlock(block);
                break;
            }
        }
    }

    /**
     * 针对 test.s 的特定优化
     * 分析样例代码中的优化机会：
     * 1. while_entry4 和其他块的跳转优化
     * 2. 多个 j ifentry7 可以考虑合并
     */
    private void optimizeTestSamplePatterns() {
        for (RISCVfunction function : module.getFunctions()) {
            // 寻找类似 test.s 中的模式
            optimizeWhileLoopPattern(function);
            optimizeIfElseChainPattern(function);
        }
    }

    /**
     * 优化while循环模式
     */
    private void optimizeWhileLoopPattern(RISCVfunction function) {
        // 寻找while循环的跳转模式并优化
        for (RISCVblock block : function.getBlocks()) {
            if (block.getName().contains("while")) {
                // 特定的while循环优化逻辑
                optimizeWhileBlock(block);
            }
        }
    }

    /**
     * 优化if-else链模式
     */
    private void optimizeIfElseChainPattern(RISCVfunction function) {
        // 寻找多个if块跳转到同一个出口的模式
        Map<RISCVblock, List<RISCVblock>> jumpTargets = new HashMap<>();

        for (RISCVblock block : function.getBlocks()) {
            RISCVinst lastInst = block.getInstructionlist().isEmpty() ?
                    null : block.getInstructionlist().getLast();

            if (lastInst instanceof RISCVj jumpInst) {
                RISCVblock target = jumpInst.getTarget();
                jumpTargets.computeIfAbsent(target, k -> new ArrayList<>()).add(block);
            }
        }

        // 如果多个块跳转到同一目标，检查是否可以优化
        for (Map.Entry<RISCVblock, List<RISCVblock>> entry : jumpTargets.entrySet()) {
            if (entry.getValue().size() > 2) {
                // 可能的优化机会
                optimizeMultipleJumpsToSameTarget(entry.getKey(), entry.getValue());
            }
        }
    }

    /**
     * 优化while块
     */
    private void optimizeWhileBlock(RISCVblock block) {
        // 实现针对while循环的特定优化
        LinkedList<RISCVinst> instructions = block.getInstructionlist();

        // 寻找可以优化的while循环模式
        for (int i = 0; i < instructions.size() - 1; i++) {
            RISCVinst inst = instructions.get(i);
            if (inst instanceof RISCVbeqz branchInst) {
                // 检查是否有优化机会
                optimizeBranchInWhile(block, branchInst, i);
            }
        }
    }

    /**
     * 优化while循环中的分支
     */
    private void optimizeBranchInWhile(RISCVblock block, RISCVbeqz branchInst, int index) {
        // 检查分支目标是否可以优化
        RISCVblock target = branchInst.getTarget();

        // 如果目标块只是简单的返回或跳转，可能可以直接内联
        if (target.getInstructionlist().size() <= 3) {
            // 考虑内联小的目标块
            considerInlining(block, branchInst, target);
        }
    }

    /**
     * 考虑内联小的基本块
     */
    private void considerInlining(RISCVblock sourceBlock, RISCVbeqz branchInst, RISCVblock targetBlock) {
        // 如果目标块很小且只有一个前驱，考虑内联
        if (hasOnlyOnePredecessor(targetBlock) && targetBlock.getInstructionlist().size() <= 2) {
            // 执行内联操作
            inlineBlock(sourceBlock, branchInst, targetBlock);
        }
    }

    /**
     * 检查块是否只有一个前驱
     */
    private boolean hasOnlyOnePredecessor(RISCVblock block) {
        int predecessorCount = 0;

        for (RISCVfunction function : module.getFunctions()) {
            for (RISCVblock otherBlock : function.getBlocks()) {
                for (RISCVinst inst : otherBlock.getInstructionlist()) {
                    if ((inst instanceof RISCVj jumpInst && jumpInst.getTarget().equals(block)) ||
                            (inst instanceof RISCVbeqz branchInst && branchInst.getTarget().equals(block))) {
                        predecessorCount++;
                        if (predecessorCount > 1) return false;
                    }
                }
            }
        }

        return predecessorCount == 1;
    }

    /**
     * 内联基本块
     */
    private void inlineBlock(RISCVblock sourceBlock, RISCVbeqz branchInst, RISCVblock targetBlock) {
        // 将目标块的指令复制到源块中
        int branchIndex = sourceBlock.getInstructionlist().indexOf(branchInst);

        // 删除原来的分支指令
        sourceBlock.deleteInstr(branchInst);

        // 插入目标块的指令
        for (int i = 0; i < targetBlock.getInstructionlist().size(); i++) {
            RISCVinst inst = targetBlock.getInstructionlist().get(i);
            sourceBlock.getInstructionlist().add(branchIndex + i, inst);
        }

        // 删除现在空的目标块
        removeBlockFromFunction(targetBlock);
    }

    /**
     * 优化多个跳转到同一目标的情况
     */
    private void optimizeMultipleJumpsToSameTarget(RISCVblock target, List<RISCVblock> sources) {
        // 分析是否可以通过重构减少跳转数量
        // 这可能涉及代码重排或条件合并

        // 简单实现：检查源块是否可以合并
        for (int i = 0; i < sources.size() - 1; i++) {
            for (int j = i + 1; j < sources.size(); j++) {
                if (canMergeBlocks(sources.get(i), sources.get(j))) {
                    // 合并这两个块
                    mergeBlockPair(sources.get(i), sources.get(j));
                    break;
                }
            }
        }
    }

    /**
     * 检查两个块是否可以合并
     */
    private boolean canMergeBlocks(RISCVblock block1, RISCVblock block2) {
        // 简单检查：如果两个块的指令序列相似
        return computeBlockSignature(block1).equals(computeBlockSignature(block2));
    }

    /**
     * 合并两个基本块
     */
    private void mergeBlockPair(RISCVblock block1, RISCVblock block2) {
        // 更新所有对block2的引用指向block1
        updateReferencesToBlock(block2, block1);

        // 删除block2
        removeBlockFromFunction(block2);
    }
}
